from cozepy import COZE_CN_BASE_URL, Coze, TokenAuth, Stream, WorkflowEvent, WorkflowEventType
import json
def run_coze_workflow1(workflow_id: str, api_token: str, user_input1:str, user_input2:str,base_url: str = COZE_CN_BASE_URL):
    """
    运行Coze工作流并处理返回的流事件

    参数:
        workflow_id: 工作流ID
        api_token: Coze API访问令牌
        base_url: Coze API基础URL(默认为中国区)
    """
    # 初始化Coze客户端
    coze = Coze(auth=TokenAuth(token=api_token), base_url=base_url)

    result_parts = []

    def handle_workflow_iterator(stream: Stream[WorkflowEvent]):
        """处理工作流事件迭代器"""
        for event in stream:
            if event.event == WorkflowEventType.MESSAGE:
                # print("收到消息:", event.message)
                result_parts.append(event.message.content)
            elif event.event == WorkflowEventType.ERROR:
                print("发生错误:", event.error)
                result_parts.append(f"ERROR: {event.error}")
            elif event.event == WorkflowEventType.INTERRUPT:
                handle_workflow_iterator(
                    coze.workflows.runs.resume(
                        workflow_id=workflow_id,
                        event_id=event.interrupt.interrupt_data.event_id,
                        resume_data="hey",
                        interrupt_type=event.interrupt.interrupt_data.type,
                    )
                )

    # 获取用户输入
    # user_input1 = input("请输入ppt主题: ")
    # user_input2 = input("请输入你想要的风格: ")

    # 运行工作流并处理结果
    handle_workflow_iterator(
        coze.workflows.runs.stream(
            workflow_id=workflow_id,
            parameters={"themes": user_input1,
                        "style": user_input2}
        )
    )
    
    return "".join(result_parts)


def run_coze_workflow2(workflow_id: str, api_token: str, outline, user_input3: int, base_url: str = COZE_CN_BASE_URL):
    """
    运行Coze工作流并处理返回的流事件，最终返回结果字符串

    参数:
        workflow_id: 工作流ID
        api_token: Coze API访问令牌
        outline: 输入的大纲内容
        user_input3: 用户选择的模板 ID
        base_url: Coze API基础URL(默认为中国区)

    返回:
        str: 收集到的消息内容或错误信息
    """
    # 初始化Coze客户端
    coze = Coze(auth=TokenAuth(token=api_token), base_url=base_url)

    result_parts = []  # 收集消息内容

    def handle_workflow_iterator(stream: Stream[WorkflowEvent]):
        """处理工作流事件迭代器"""
        for event in stream:
            if event.event == WorkflowEventType.MESSAGE:
                result_parts.append(event.message.content)
            elif event.event == WorkflowEventType.ERROR:
                result_parts.append(f"ERROR: {event.error}")
            elif event.event == WorkflowEventType.INTERRUPT:
                handle_workflow_iterator(
                    coze.workflows.runs.resume(
                        workflow_id=workflow_id,
                        event_id=event.interrupt.interrupt_data.event_id,
                        resume_data="hey",
                        interrupt_type=event.interrupt.interrupt_data.type,
                    )
                )

    # 运行工作流并处理结果
    handle_workflow_iterator(
        coze.workflows.runs.stream(
            workflow_id=workflow_id,
            parameters={"input_outline": outline, "input_themeId": user_input3}
        )
    )

    return ''.join(result_parts)  # 合并所有片段并返回
def device(result: str):
    """
    用于处理卡面一部获取的结果
    """
    data = json.loads(result)
    outline = data.get("outline")
    themes = data.get("themes")

    return outline, themes

# 使用示例
if __name__ == "__main__":
    # 配置参数
    WORKFLOW_ID1 = "7524165514211541002"
    WORKFLOW_ID2 = "7524182748136177664"
    API_TOKEN = "pat_uQoTyqb5Mzce6Rbz9DnqtyAy8vndHkjxh6pFAzsmy8QvF5aHPvm9dk50r0sSMYQN"

    # 运行工作流
    first_result = run_coze_workflow1(workflow_id=WORKFLOW_ID1, api_token=API_TOKEN, user_input1="数学", user_input2="山水风")
    # print("以下是ppt主题和模板",first_result)
    outline=json.loads(first_result).get("outline")
    themes=json.loads(first_result).get("theme")
    # outline, themes = device(first_result)
    print(first_result)
    print("--------------")
    print("以下是ppt主题",outline)
    print("--------------")
    print("以下是ppt模板",themes)
    a=input("请输入你想使用的模板: ")
    b=run_coze_workflow2(workflow_id=WORKFLOW_ID2, api_token=API_TOKEN, outline=outline, user_input3=a)
    b=json.loads(b).get("output")
    print(b)